# !/usr/bin/env/ python3.7
# -*- coding:utf-8 -*-
# optimize_chatbot.py
# optimize chat robot 
# author: luoxiang


import pickle
import time


class RobotUI:
    def __init__(self, name, admin):
        self.name = name
        self.admin = admin
        
        
    def talk(self):
        words = input(f'{self.user}：')
        answer = self.admin.answer(words)
        print(f'{self.name}: {answer}')
    
    def welcome(self):
        print(f'{self.name}:这个名字我很喜欢(*^▽^*)')
        time.sleep(1)
        self.user = input(f'{self.name}:怎么称呼您？(请输入姓名：)')
        print(f'{self.name}:很高兴能和你聊天，想聊点什么？')


class RobotAdmin:
    def __init__(self, name):
        self.data = self.load_data()
        self.name = name


    def load_data(self):
        with open('corpus.pkl', 'rb') as f:
            data = pickle.loads(f.read())
            print(data)
            return data

    def save_data(self):
        with open('corpus.pkl', 'wb') as f:
            f.write(pickle.dumps(self.data))


    @staticmethod
    def add_corpus(self):
        key = input(f'{self.name}:请输入问题：')
        value = input(f'{self.name}:请输入答案：')
        self.data[key] = value
        return self.data

        

    def answer(self, listen):
        ''' answer something from listen'''
        if listen in self.data:
            return self.data[listen]
        elif listen not in self.data:
            print(f'{self.name}:我不太明白你的意思!')
            time.sleep(1)
            choice = input(f'{self.name}:你能帮我完善语料库吗？愿意：y，不愿意：q')
            if choice == 'y': 
                while True:
                    print(self.data)
                    RobotAdmin.add_corpus(self)
                    choose = input(f"{self.name}:还想继续添加吗？愿意：y，结束：n")
                    if choose == 'y':
                        continue
                    elif choose == 'n':
                        return self.data['n']

            elif choice == 'q':
                return self.data['q']
            
        

def main():
    """
    请扩展如下功能，欢迎写好了到群里炫富！

    多轮交互，还记得while循环吧，记得给个退出提示，刹车很重要！
    根据关键词进行回答，前面学的字符串查找和正则不是吃素的哦！
    扩展admin功能，可以不生成对象情况下自定义添加数据（staticmethod用过没？）
    你还想干啥？Just Do IT！
    """
    
    print('我是聊天机器人！')
    time.sleep(1)
    name = input('能帮我起个名字吗？：')
    admin = RobotAdmin(name)
    robot = RobotUI(name, admin)
    robot.welcome()
    add = True
    while(add):
        robot.talk()
        add = input(f"{name}:还想继续和我聊天吗？继续：按任意键，退出：no")
        if add == 'no':
            print(f'{name}:很高兴认识你，有空又聊！')
            admin.save_data()
            break
        else:
            continue
    


if __name__ == '__main__':
    main()
